Skip to content

fix: pass total_tokens for SSM seq_idx in packed sequences - #3484

Merged
cuichenx merged 1 commit into
mainfrom
fix/pass-total-tokens-for-ssm-seq-idx
Apr 25, 2026
Merged

fix: pass total_tokens for SSM seq_idx in packed sequences#3484
cuichenx merged 1 commit into
mainfrom
fix/pass-total-tokens-for-ssm-seq-idx

Conversation

@yaoyu-33

@yaoyu-33 yaoyu-33 commented Apr 22, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Pass total_tokens through get_packed_seq_params() to PackedSeqParams, enabling seq_idx generation for hybrid SSM/Mamba models during varlen (packed sequence) training
  • Without this, seq_idx remains None and SSM state bleeds across sequence boundaries in packed batches, corrupting training
  • Handles pipeline parallelism correctly: uses tokens.size(1) on the first PP stage and labels.size(1) on the last PP stage (matching upstream Megatron-LM pretrain_mamba.py)
  • Affects all hybrid SSM models (Nemotron Nano, NemotronH, Nemotron Nano V2, Nemotron 3 Super), not just Nemotron Nano
  • Backward compatible — total_tokens defaults to None for non-SSM models

Fixes #3474

Test plan

  • Unit tests added for get_packed_seq_params covering total_tokens → seq_idx generation, backward compat (seq_idx=None without total_tokens), and the cu_seqlens_unpadded path
  • Verify Nemotron Nano SFT with sequence packing no longer produces corrupted SSM states
  • Verify non-SSM models (e.g. Llama, Qwen) are unaffected (no seq_idx generated when total_tokens absent)

🤖 Generated with Claude Code

@copy-pr-bot

copy-pr-bot Bot commented Apr 22, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@yaoyu-33 yaoyu-33 mentioned this pull request Apr 22, 2026
Closed
@yaoyu-33 @claude
fix: pass total_tokens to PackedSeqParams for SSM seq_idx generation
5e50ac2 
Without total_tokens, PackedSeqParams.__post_init__ skips seq_idx
computation, causing Mamba/SSM state to bleed across sequence boundaries
in varlen (packed sequence) training. This affects all hybrid SSM models
(Nemotron Nano, NemotronH, Nemotron Nano V2, Nemotron 3 Super), not
just Nemotron Nano.

Fixes #3474

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
@yaoyu-33
yaoyu-33 force-pushed the fix/pass-total-tokens-for-ssm-seq-idx branch from 2a2e10c to 5e50ac2 Compare April 22, 2026 21:44
cuichenx
cuichenx approved these changes Apr 23, 2026
View reviewed changes

@cuichenx cuichenx left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@yaoyu-33 yaoyu-33 added bug Something isn't working area:training Training loop, callbacks, and runtime integration needs-review PR is ready for code review and waiting on a reviewer labels Apr 23, 2026
@cuichenx cuichenx mentioned this pull request Apr 23, 2026
Closed
5 tasks
@cuichenx cuichenx added ready-to-merge PR is approved, current, and only waiting for CI to pass before merge and removed needs-review PR is ready for code review and waiting on a reviewer labels Apr 23, 2026
@cuichenx

cuichenx commented Apr 23, 2026

Copy link
Copy Markdown
Contributor

/ok to test 5e50ac2

@cuichenx
cuichenx merged commit 10d6e76 into main Apr 25, 2026
145 of 147 checks passed
@cuichenx
cuichenx deleted the fix/pass-total-tokens-for-ssm-seq-idx branch April 25, 2026 00:59
@pruprakash

pruprakash commented Apr 28, 2026

Copy link
Copy Markdown
Contributor

QA RCCA Analysis

1. Fix Reference

2. Root Cause

total_tokens was not passed through get_packed_seq_params() to PackedSeqParams, preventing seq_idx generation for hybrid SSM/Mamba models during varlen (packed sequence) training. Without this, seq_idx remains None and SSM state bleeds across sequence boundaries in packed batches, corrupting training.

3. Trigger Configuration

  • Training hybrid SSM models (Nemotron Nano, NemotronH, Nemotron Nano V2, Nemotron 3 Super)
  • Using packed sequences / varlen training

4. Nature of the Bug

Classification: CODE BUG - SSM state corruption in packed sequence training

5. Existing Test Coverage

In Fix PR: YES - 1 test file modified:

  • tests/unit_tests/training/utils/test_packed_seq_utils.py

In NMFW Tests: Not specifically covering SSM packed sequence scenarios

6. Coverage Assessment

Test Type Exists Covers Bug
Fix PR unit tests YES YES
NMFW regression tests NO N/A

7. New Regression Test

NOT NEEDED - Fix PR includes unit tests for get_packed_seq_params covering total_tokens parameter.

8. Conclusion

Verdict: ADEQUATE COVERAGE - Fix PR includes unit tests verifying total_tokens handling for SSM models.

@cuichenx cuichenx mentioned this pull request Apr 29, 2026
Merged
3 tasks
vasunvidia pushed a commit to vasunvidia/Megatron-Bridge that referenced this pull request Jun 10, 2026
@yaoyu-33 @claude @vasunvidia
fix: pass total_tokens for SSM seq_idx in packed sequences (NVIDIA-Ne…
aa3a69c 
…Mo#3484)

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@cuichenx cuichenx cuichenx approved these changes

Assignees

No one assigned

Labels

area:training Training loop, callbacks, and runtime integration bug Something isn't working qa_rcca_done ready-to-merge PR is approved, current, and only waiting for CI to pass before merge

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[bug] nemo3nano varlen missing totoal_tokens for ssm layers

3 participants

@yaoyu-33 @cuichenx @pruprakash